#ifndef ROBOTICS_H
#define ROBOTICS_H

#include <math.h>
#include <stdio.h>

typedef struct {
    float data[4][4];
} Matrix4x4;

typedef struct {
    float x, y, z;
    float roll, pitch, yaw;
} Pose;

Matrix4x4 matrix_mult(Matrix4x4 a, Matrix4x4 b) {
    Matrix4x4 result;
    for (int i = 0; i < 4; i++) {
        for (int j = 0; j < 4; j++) {
            result.data[i][j] = 0;
            for (int k = 0; k < 4; k++) {
                result.data[i][j] += a.data[i][k] * b.data[k][j];
            }
        }
    }
    return result;
}

Matrix4x4 joint_transform(float theta, float d, float a, float alpha) {
    Matrix4x4 T;
    float ct = cosf(theta);
    float st = sinf(theta);
    float ca = cosf(alpha);
    float sa = sinf(alpha);

    T.data[0][0] = ct;
    T.data[0][1] = -st * ca;
    T.data[0][2] = st * sa;
    T.data[0][3] = a * ct;

    T.data[1][0] = st;
    T.data[1][1] = ct * ca;
    T.data[1][2] = -ct * sa;
    T.data[1][3] = a * st;

    T.data[2][0] = 0;
    T.data[2][1] = sa;
    T.data[2][2] = ca;
    T.data[2][3] = d;

    T.data[3][0] = T.data[3][1] = T.data[3][2] = 0;
    T.data[3][3] = 1;
    return T;
}

Pose fkine(float value[7]) {
    // 关节参数计算
    Matrix4x4 T[7];
    
    // 关节0: Prismatic (d = value[0])
    T[0] = joint_transform(0, value[0], 0, -M_PI_2);
    
    // 关节1: Revolute (theta = value[1])
    T[1] = joint_transform(value[1], 0, 0, -M_PI_2);
    
    // 关节2: Revolute (theta = value[2])
    T[2] = joint_transform(value[2], 0, 325.9, M_PI);
    
    // 关节3: Revolute (theta = value[3])
    T[3] = joint_transform(value[3], 59.6, 239.78, M_PI);
    
    // 关节4: Revolute (theta = value[4] + pi/2)
    T[4] = joint_transform(value[4] + M_PI_2, 59.53, 0, M_PI_2);
    
    // 关节5: Revolute (theta = value[5] + pi/2)
    T[5] = joint_transform(value[5] - M_PI_2, 66.14, 0, M_PI_2);
    
    // 关节6: Revolute (theta = value[6])
    T[6] = joint_transform(value[6], 104.5, 0, 0);

    // 计算总变换矩阵
    Matrix4x4 T_total = T[0];
    for (int i = 1; i < 7; i++) {
        T_total = matrix_mult(T_total, T[i]);
    }

    // 提取位置
    Pose pose;
    pose.x = T_total.data[0][3];
    pose.y = T_total.data[1][3];
    pose.z = T_total.data[2][3];

    // 提取欧拉角 (Z-Y-X顺序，yaw-pitch-roll)
    float r31 = T_total.data[2][0];
    float r11 = T_total.data[0][0];
    float r21 = T_total.data[1][0];
    float r32 = T_total.data[2][1];
    float r33 = T_total.data[2][2];

    pose.pitch = asinf(-r31);
    float cp = cosf(pose.pitch);

    if (fabsf(cp) > 1e-6) {
        pose.roll = atan2f(r32 / cp, r33 / cp);
        pose.yaw = atan2f(r21 / cp, r11 / cp);
    } else {
        // 万向锁处理
        pose.roll = 0;
        pose.yaw = atan2f(-T_total.data[0][1], T_total.data[1][1]);
    }

    return pose;
}

int main(){
    // float joint_values[7]={0,0,0,0,0,0,0};
    float joint_values[7]={0,0.1,0.2,0.3,0.4,0.5,0.6};
    Pose output=fkine(joint_values);
    printf("pose :  %f,%f,%f,%f,%f,%f \n", output.x, output.y, output.z, output.roll, output.pitch, output.yaw);
    return 0;
}

#endif // ROBOTICS_H